#' ----------------------------------------------------------------------------
#' Different Specifications of BAM Models
#' 
#' Description: This script estimates various Bayesian Aldrich-McKelvey models
#' using data prepared by the 00_data.R script.

# Install and load required packages -------------------------------------------
if (!require("pacman")) install.packages("pacman")
p_load(tidyverse, rstan, hbamr)

# Load Data --------------------------------------------------------------------
# Wide dataset
df_wide <- readRDS("outputs/df_global_expert.rds")

# BAM MODEL --------------------------------------------------------------------

# Data to long format: only Left-right placements and anchoring vignettes
df_long_raw <- df_wide %>%
  ungroup() %>%
  mutate(obs = row_number()) %>%
  rename(lrecon_002 = party_center,
         lrecon_000 = party_left,
         lrecon_001 = party_right) %>%
  pivot_longer(contains("lrecon"), names_to = "party", values_to = "place") %>%
  group_by(party) %>%
  mutate(pnum = cur_group_id()) %>%
  ungroup()

# Data to long without NAs (sparse dataset) and standardized placements
df_long <- df_long_raw %>%
  select(region_id, country_id, country, expert_id, obs, party, pnum, place) %>%
  na.omit() %>% # delete all missing
  mutate(place = (place - mean(place)) / sd(place)) # standardize

# Save data (will be useful later for the analysis)
saveRDS(df_long, 'outputs/df_bam.rds')

# Data to Stan
stan_data <- list(J = max(df_long$obs),
                  N = max(df_long$pnum),
                  nobs = nrow(df_long),
                  ctry = df_long$country_id,
                  pty = df_long$pnum,
                  id = df_long$obs,
                  place = df_long$place)

# Estimate model
bam_fit <- stan("scripts/stan/00_BAM.stan",
                data = stan_data,
                seed = 100,
                chains = 4,
                iter = 4000,
                warmup = 2000,
                control = list(adapt_delta = 0.99,
                               max_treedepth = 15),
                cores = 4,
                pars = c("Z", "beta", "alpha", "sigma",
                         "sigma_j", "sigma_n"))
saveRDS(bam_fit, 'outputs/bam.rds')

# HBAM MODELS ------------------------------------------------------------------
# Estimate HBAM
hbam_fit <- hbam(df_wide$expert_lrgen, # self-placements
                 df_wide[, which(names(df_wide) == "party_left"):ncol(df_wide)], # Stimuli
                 req_valid = 3, # 3 valid placements
                 req_unique = 2, # 2 unique placements
                 model = "HBAM_NF", # No flipping
                 cores = 4,
                 chains = 4,
                 warmup = 2000,
                 iter = 8000,
                 seed = 100,
                 control = list(adapt_delta = 0.999))
saveRDS(hbam_fit, 'outputs/hbam_base.rds')

# Estimate Hbam Multi
hbam_multi_fit <- hbam(df_wide$expert_lrgen, # self-placements
                       df_wide[, which(names(df_wide) == "party_left"):ncol(df_wide)], # Stimuli
                       model = "HBAM_MULTI_NF",
                       group_id = as.numeric(as.factor(df_wide$country)),
                       req_valid = 3,
                       req_unique = 2,
                       cores = 4,
                       chains = 4,
                       warmup = 2000,
                       iter = 8000,
                       seed = 100,
                       control = list(adapt_delta = 0.99,
                                      max_treedepth = 12))
saveRDS(hbam_multi_fit, 'outputs/hbam_multi.rds')

# Save session information for reproducibility
writeLines(capture.output(sessionInfo()), "outputs/session_info/01_BAM_models.txt")